//
// Created by mio on 25-5-13.
//

#include "H264Encoder.h"
#include "VideoConvert.h"

#include <iomanip>
#include <iostream>

extern "C" {
#include <libavutil//rational.h>
#include <libavutil/error.h>
#include <libavutil/opt.h>
}

H264Encoder::H264Encoder() : pts_(0), width_(0), height_(0), force_idr_(false), rgba_frame_(nullptr),
    h264_packet_(nullptr), converter_(nullptr), current_codec_name_("")
{
    rgba_frame_.reset(av_frame_alloc(), [](AVFrame* frame) { av_frame_free(&frame); });
    h264_packet_.reset(av_packet_alloc(), [](AVPacket* packet) { av_packet_free(&packet); });
}

H264Encoder::~H264Encoder()
{
    Close();
}

bool H264Encoder::Open(AVConfig &video_config)
{
    if (is_initialzed_) return false;

    config_ = video_config;

    const char* encoders[][2] = {
        {"h264_nvenc", "NVIDIA NVENC"}, // NVIDIA
        {"h264_amf", "AMD AMF"},       // AMD
        {"h264_qsv", "Intel QSV"},     // Intel Quick Sync Video
        {nullptr, nullptr} // 列表结束标记
    };

    AVCodec* found_codec = nullptr;
    AVDictionary* options = nullptr;


    for (int i = 0; encoders[i][0] != nullptr; i++) {
        found_codec = (AVCodec*)avcodec_find_encoder_by_name(encoders[i][0]);
        if (found_codec) {
            codecContext_ = avcodec_alloc_context3(found_codec);
            if (!codecContext_) { found_codec = nullptr; continue; }
            std::string current_hw_encoder_name_str = encoders[i][0];

            codecContext_->width = video_config.video.width;
            codecContext_->height = video_config.video.height;
            codecContext_->framerate = { (qint32)video_config.video.framerate, 1 };
            codecContext_->time_base = { 1, (qint32)video_config.video.framerate };
            // 硬件编码器通常偏好 AV_PIX_FMT_NV12，但 YUV420P 通常也能被接受或自动转换
            // 如果特定硬件编码器强制要求NV12等，这里可能需要调整，
            // 或者让VideoConvert输出对应的格式
            codecContext_->pix_fmt = AV_PIX_FMT_YUV420P;
            codecContext_->gop_size = 30;
            codecContext_->max_b_frames = 0;
            codecContext_->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

            // 设置硬件特定的低延迟选项
            if (current_hw_encoder_name_str == "h264_nvenc") {
                av_opt_set(codecContext_->priv_data, "preset", "llhp", 0); // Low-Latency
                av_opt_set(codecContext_->priv_data, "tune", "ll", 0); // // Low-Latency tune
                av_opt_set_int(codecContext_->priv_data, "zerolatency", 1, 0); // // 确保零延迟模式
                av_opt_set_int(codecContext_->priv_data, "delay", 0, 0);          // 强制无内部缓冲延迟
                // av_opt_set(codecContext_->priv_data, "profile", "baseline", 0); // Baseline profile for widest compatibility in low latency
            }
            else if (current_hw_encoder_name_str == "h264_amf") {
                av_opt_set(codecContext_->priv_data, "usage", "ultralowlatency", 0);
                av_opt_set(codecContext_->priv_data, "quality", "speed", 0);
            }
            else if (current_hw_encoder_name_str == "h264_qsv") {
                // 对于 QSV, preset 通常指整体质量/速度权衡
                av_opt_set(codecContext_->priv_data, "preset", "ultrafast", 0); // 或者 "veryfast"
                av_opt_set_int(codecContext_->priv_data, "async_depth", 1, 0);   // 关键: QSV异步深度为1是低延迟的核心
                av_opt_set_int(codecContext_->priv_data, "look_ahead_depth", 0, 0); // 禁用先行侦测
                // QSV的low_power选项可能影响延迟 "on" 或 "off"
                // av_opt_set(codecContext_->priv_data, "low_power", "on", 0); // 某些场景下可能有用，但需测试
            }

            if (avcodec_open2(codecContext_, found_codec, nullptr) == 0)
            {
                current_codec_name_ = encoders[i][1];
                codec_ = found_codec;
                qDebug() << "H264Encoder: Using hardware encoder: " << current_codec_name_;
                break;
            }
            else {
                avcodec_free_context(&codecContext_);
                codecContext_ = nullptr;
                found_codec = nullptr;
            }
        }
    }

    if (!codec_) {
        codec_ = (AVCodec*)avcodec_find_encoder(AV_CODEC_ID_H264);
        if (!codec_) { Close(); return false; }

        codecContext_ = avcodec_alloc_context3(codec_);
        if (!codecContext_) { Close(); return false; }

        codecContext_->width = video_config.video.width;
        codecContext_->height = video_config.video.height;
        codecContext_->framerate = {(qint32)video_config.video.framerate, 1 };
        codecContext_->time_base = { 1, (qint32)video_config.video.framerate };
        codecContext_->pix_fmt = AV_PIX_FMT_YUV420P;
        codecContext_->gop_size = 30;
        codecContext_->max_b_frames = 0;
        codecContext_->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

        // 加上参数，降低编码延迟
        codecContext_->rc_min_rate = config_.video.bitrate;
        codecContext_->rc_max_rate = config_.video.bitrate;
        codecContext_->rc_buffer_size = config_.video.bitrate;

        // 设置字典 libx264 特定参数
        av_opt_set(codecContext_->priv_data, "tune", "zerolatency", 0); // 极速编码
        av_opt_set(codecContext_->priv_data, "preset", "ultrafast", 0); // 极速编码

        if (avcodec_open2(codecContext_, codec_, nullptr) != 0) {
            Close();
            return false;
        }

        codec_ = found_codec;
        current_codec_name_ = "libx264 (Software)";
        qDebug() << "H264Encoder: Using software encoder: " << current_codec_name_;
    }

    width_ = config_.video.width;
    height_ = config_.video.height;
    is_initialzed_ = true;
    return true;
}

void H264Encoder::Close()
{
    if (is_initialzed_) {
        width_ = 0;
        height_ = 0;
        pts_ = 0;
        is_initialzed_ = false;
        if (codecContext_) {
            avcodec_free_context(&codecContext_);
            codecContext_ = nullptr;
        }
        if (converter_) {
            converter_->Close();
            converter_.reset();
            converter_ = nullptr;
        }
    }
}

AVPacketPtr H264Encoder::Encode(const quint8 *data, quint32 width, quint32 height, quint32 data_size, quint64 pts)
{
    if(!is_initialzed_) { return nullptr; }

    // auto startEncode = std::chrono::high_resolution_clock::now();
    if (width_ != width || height_ != height || !converter_)
    {
        converter_.reset(new VideoConvert());
        if (!converter_->Open(width, height, config_.video.format,
                            codecContext_->width, codecContext_->height, codecContext_->pix_fmt))
        {
            converter_->Close();
            converter_.reset();
            converter_ = nullptr;
        }

        rgba_frame_->width = width_;
        rgba_frame_->height = height_;
        rgba_frame_->format = config_.video.format;

        if (av_frame_get_buffer(rgba_frame_.get(), 32) != 0) { return nullptr; }
    }

    memcpy(rgba_frame_->data[0], data, data_size);

    if (!converter_) return nullptr;

    AVFramePtr out_frame = nullptr;
    if (converter_->Convert(rgba_frame_, out_frame) <= 0) { return nullptr; }

    out_frame->pts = pts >= 0 ? pts : pts_++;

    if (avcodec_send_frame(codecContext_, out_frame.get()) < 0) { return nullptr; }

    int ret = avcodec_receive_packet(codecContext_, h264_packet_.get());
    if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) { return nullptr; }
    else if (ret < 0) { return nullptr; }

    // auto end = std::chrono::high_resolution_clock::now();
    // auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end - startEncode);
    // qDebug() << "Encode Duration: " << duration.count() << " microseconds";

    return h264_packet_;
}
